#!/usr/bin/env python3
# pur_policy_client.py
# 远程策略客户端：连接pur_policy服务器，加载DRL模型，接收状态并返回计算结果
import socket
import struct
import numpy as np
import os
import logging
import time
from drl_policy_driver import DRLPolicy

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def receive_all(sock, length):
    """可靠地接收指定长度的数据"""
    data = b""
    while len(data) < length:
        packet = sock.recv(length - len(data))
        if not packet:
            raise ConnectionError("连接中断")
        data += packet
    return data


def connect_to_server(server_ip, port):
    """连接到策略服务器"""
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(10)  # 新增：10s 读超时
    sock.connect((server_ip, port))
    return sock


def initialize_session(sock):
    """初始化会话，接收初始图结构更新包"""
    # 接收包头: 期望 0x01 (图结构更新)
    cmd = sock.recv(1)
    if not cmd or cmd != b"\x01":
        raise ValueError(f"初始包头无效: {cmd}")
    # 接收节点数
    node_num_data = receive_all(sock, 4)
    node_num = struct.unpack(">I", node_num_data)[0]
    logger.info(f"接收到节点数: {node_num}")

    # 接收邻接矩阵长度
    len_data = receive_all(sock, 4)
    len_adj = struct.unpack(">I", len_data)[0]
    logger.info(f"接收到邻接矩阵长度: {len_adj}")

    # 接收邻接矩阵字节
    adj_bytes = receive_all(sock, len_adj)
    logger.info(f"接收到邻接矩阵字节,大小为 {len(adj_bytes)}")
    adj_matrix = unpack_adj_matrix(adj_bytes, node_num)
    logger.info("接收到初始邻接矩阵")

    return adj_matrix


def handle_requests(sock, drl_policy):
    """处理请求循环，支持图结构更新和策略请求"""
    while True:
        try:
            # 接收包头
            cmd = sock.recv(1)
            if not cmd:
                raise ConnectionError("连接中断")
            if cmd == b"\x01":  # 图结构更新
                # 接收节点数
                node_num_data = receive_all(sock, 4)
                node_num = struct.unpack(">I", node_num_data)[0]
                logger.info(f"接收到图更新节点数: {node_num}")
                # 接收邻接矩阵长度
                len_data = receive_all(sock, 4)
                len_adj = struct.unpack(">I", len_data)[0]
                logger.info(f"接收到图更新邻接矩阵长度: {len_adj}")
                # 接收邻接矩阵字节
                adj_bytes = receive_all(sock, len_adj)
                adj_matrix = unpack_adj_matrix(adj_bytes, node_num)
                logger.info("接收到图结构更新")
                # 更新DRL策略的图
                drl_policy.update_graph(adj_matrix)
                logger.info("DRL策略图结构已更新")

            elif cmd == b"\x02":  # 策略解算请求
                # 接收数据段: 当前节点状态
                state_data = receive_all(sock, 12)
                evader_node, p1_node, p2_node = struct.unpack(">III", state_data)
                logger.info(
                    f"<<< 接收状态: 逃避者={evader_node}, 追捕者1={p1_node}, 追捕者2={p2_node}"
                )

                # 调用DRL计算
                next_evader = evader_node
                next_p1, next_p2 = drl_policy.calculate_next_move(
                    evader_node, p1_node, p2_node
                )

                # 发送响应: 包头 0x03 + 数据段
                response = b"\x03" + struct.pack(">III", next_evader, next_p1, next_p2)
                sock.sendall(response)
                logger.info(
                    f">>> 发送策略: 逃避者={next_evader}, 追捕者1={next_p1}, 追捕者2={next_p2}"
                )
            else:
                logger.warning(f"未知包头: {cmd}")
        except socket.timeout:
            logger.warning("接收超时，跳过")
            continue
        except ConnectionError as e:
            logger.warning(f"连接错误: {e}")
            raise  # 抛出以触发重连


def unpack_adj_matrix(packed_bytes, node_num):
    """从bit-packed bytes解包成邻接矩阵"""
    total_bits = node_num * node_num
    num_bytes = len(packed_bytes)
    flat_bits = np.zeros(total_bits, dtype=bool)
    for i in range(total_bits):
        byte_idx = i // 8
        bit_idx = 7 - (i % 8)  # MSB first
        if byte_idx < num_bytes and (packed_bytes[byte_idx] & (1 << bit_idx)):
            flat_bits[i] = True
    return flat_bits.reshape((node_num, node_num)).astype(np.uint8)


def main():
    server_ip = "127.0.0.1"
    file_dir_path = os.path.dirname(__file__)
    model_path = os.path.join(file_dir_path, "./pursuer/model/pursuer_model/policy.pth")
    port = 6002

    drl_policy = None
    adj_matrix = None

    while True:
        sock = None
        try:
            # 连接到策略服务器
            logger.info(f"正在连接策略服务器  {server_ip}:{port}")
            sock = connect_to_server(server_ip, port)
            logger.info(f"已连接到策略服务器 {server_ip}:{port}")

            # 初始化会话和模型/图更新
            adj_matrix = initialize_session(sock)
            if drl_policy is None:
                # 首次连接，加载DRL模型
                logger.info("正在加载DRL模型...")
                drl_policy = DRLPolicy(adj_matrix, model_path)
                logger.info(f"DRL模型已加载: {model_path}")
            else:
                # 重新连接，只更新图相关部分
                logger.info("重新连接成功，更新图结构...")
                drl_policy.update_graph(adj_matrix)
                logger.info("图结构已更新，使用现有模型")

            logger.info("策略客户端就绪，准备处理请求...")

            # 处理请求
            handle_requests(sock, drl_policy)

        except ConnectionError as e:
            logger.warning(f"连接错误: {e}")
            if sock:
                sock.close()
            logger.info("等待5秒后尝试重新连接...")
            time.sleep(5)
            continue

        except Exception as e:
            logger.error(f"发生未预期错误: {e}")
            if sock:
                sock.close()
            logger.info("等待5秒后尝试重新连接...")
            time.sleep(5)
            continue

        finally:
            if sock:
                sock.close()
                logger.info("客户端连接已关闭")


if __name__ == "__main__":
    main()
